# 该文件通过使用DeepMTT模型预测轨迹
import time
import tensorflow as tf
import numpy as np
import os
from tensorflow.keras import models
from custome_activations import noisy_activation
from batchdata_derive3 import create_all_data
from DeepMTT_build import maxout_activation
import matplotlib.pyplot as plt


# 准备输入数据
def pre_process():
    true_traj, obs_traj, est_traj, error_between = create_all_data(batch_size=64,
                                                                   data_len=50,
                                                                   ini_pos_noise=30,
                                                                   ini_vel_noise=3,
                                                                   batch_num=10)
    weights = np.max(np.abs(est_traj), axis=(1, 2))
    weights = weights.reshape(-1, 1, 1)  # 重塑为[batch_size, 1, 1]的形状

    # 标准化所有轨迹
    processed_inputs = est_traj / weights
    return processed_inputs, true_traj, est_traj

def predict_data():
    # 加载模型
    my_model = models.load_model('/root/lanyun-tmp/MTT419/tmp/ckpt/DeepMTT_086-455.93-425.00.keras',
                                        custom_objects={'noisy_activation': noisy_activation,
                                                        'maxout_activation': maxout_activation})
    print('模型加载成功')

    # 准备输入数据
    my_inputs, true_orbit, est_traj = pre_process()
    # 批量预测
    prediction = my_model.predict(my_inputs)
    print("预测原始结果：", prediction)
    # 根据神经网络输出还原轨迹
    predict_orbit = prediction + est_traj
    print(f"predict_orbit形状：{predict_orbit.shape}")
    picture_folder = '/root/lanyun-tmp/MTT419/tmp/picture'
    os.makedirs(picture_folder, exist_ok=True)  # 创建文件夹
    # 画图
    for i in range(20):
        plt.figure()   # 创建图表1
        plt.plot(true_orbit[i, :, 0], true_orbit[i, :, 1], linestyle='-')
        plt.plot(predict_orbit[i, :, 0], predict_orbit[i, :, 1], linestyle='--')
        plt.plot(est_traj[i, :, 0], est_traj[i, :, 1], linestyle='-.')
        plt.title('真实轨迹与DeepMTT网络预测轨迹')  # 标题
        plt.xlabel('X')  # X轴标签
        plt.ylabel('Y')  # Y轴标签
        plt.legend()  # 显示图例
        plt.grid(True)  # 显示网格
        plt.show()  # 显示图形
        filename = f'trajectory_{i+1:02d}.png'
        full_path = os.path.join(picture_folder, filename)
        plt.savefig(full_path, dpi=300, bbox_inches='tight')
        plt.close()
        time.sleep(0.5)
    print("\n循环生成图片完成")

if __name__ == '__main__':
    predict_data()
